import os
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import DashScopeEmbeddings
import difflib

class RetrievalEvaluator:
    def __init__(self, index_dir=None, k=3, threshold=0.4):
        # 默认向量库路径
        if index_dir is None:
            index_dir = r"D:\Work\Python AI\P5\P5Code\faiss_index"
        self.k = k
        self.threshold = threshold
        # 加载 embeddings（需与构建向量库时一致）
        from dotenv import load_dotenv
        load_dotenv()
        dashscope_api_key = os.getenv("DASHSCOPE_API_KEY")
        self.embeddings = DashScopeEmbeddings(
            dashscope_api_key=dashscope_api_key,
            model="text-embedding-v1"
        )
        # 加载本地向量库
        self.vectorstore = FAISS.load_local(index_dir, self.embeddings, allow_dangerous_deserialization=True)

    def is_similar(self, a, b):
        """判断两个字符串相似度是否超过阈值"""
        return difflib.SequenceMatcher(None, a, b).ratio() > self.threshold

    def compute_mrr(self, test_data):
        """
        test_data: List[Dict], 每个dict包含 'question' 和 'gold_answer'
        """
        reciprocal_ranks = []
        for item in test_data:
            if "question" not in item or "gold_answer" not in item:
                print(f"样本缺少字段: {item}")
                continue
            question = item["question"]
            gold_answer = item["gold_answer"]
            retriever = self.vectorstore.as_retriever(search_kwargs={"k": self.k})
            docs = retriever.invoke(question)
            found = False
            for idx, doc in enumerate(docs):
                doc_answer = doc.page_content
                print(f"Q: {question} | Doc{idx+1}: {doc_answer}")
                print(f"相似度: {difflib.SequenceMatcher(None, gold_answer.strip(), doc_answer).ratio():.3f}")
                if gold_answer.strip() in doc_answer or self.is_similar(gold_answer.strip(), doc_answer):
                    reciprocal_ranks.append(1.0 / (idx + 1))
                    found = True
                    break
            if not found:
                reciprocal_ranks.append(0.0)
        mrr = sum(reciprocal_ranks) / len(reciprocal_ranks) if reciprocal_ranks else 0.0
        return mrr